'''
    保存加载的 net.pkl 其实一个字典，通常包含如下内容
    1.网络结构：输入尺寸、输出尺寸以及隐藏层信息，以便能够在加载时重建模型。
    2.模型的权重参数：包含各网络层训练后的可学习参数，
    可以在模型实例上调用 state_dict() 方法来获取，
    比如前面介绍只保存模型权重参数时用到的 model.state_dict()。
    3.优化器参数：有时保存模型的参数需要稍后接着训练，那么就必须保存优化器的状态和所其使用的超参数，
    也是在优化器实例上调用 state_dict() 方法来获取这些参数。
    4.其他信息：有时我们需要保存一些其他的信息，比如 epoch，batch_size 等超参数。
'''
import torchvision
import torch


epoch = 100

model= torchvision.models.vgg16()

optiizer = torch.optim.Adam(model.parameters(),lr=0.01)

# saving a checkpoint assuming the network class named ClassNet
checkpoint = {
    'model':torchvision.models.vgg16(),
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optiizer.state_dict(),
    'epoch': epoch
}
# checkpoint 是个字典，里面有4个键值对，分别表示网络模型的不同信息。
# 加载上面保存的自定义的模型：
def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = checkpoint['model']  # 提取网络结构
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载网络权重参数
    optimizer = torch.optim.Adam(model.parameters(),lr=0.01)
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # 加载优化器参数
    for parameter in model.parameters():
        parameter.requires_grad = False
    model.eval()

    return model

model = load_checkpoint("checkpoint.pkl")
# 载模型只是为了进行推理测试，则将每一层的 requires_grad 置为 False，即固定这些权重参数；
# 还需要调用 model.eval() 将模型置为测试模式，主要是将 dropout 和 batch normalization 层进行固定，否则模型的预测结果每次都会不同。
# 希望继续训练，则调用 model.train()，以确保网络模型处于训练模式。
# state_dictdwd() dwqdwqdnwqojdqw也是一个Python字典对象，model.state_dict() 将每一层的可学习参数映射为参数矩阵，其中只包含具有可学习参数的层(卷积层、全连接层等)。